/*

 * Licensed to the Apache Software Foundation (ASF) under one

 * or more contributor license agreements.  See the NOTICE file

 * distributed with this work for additional information

 * regarding copyright ownership.  The ASF licenses this file

 * to you under the Apache License, Version 2.0 (the

 * "License"); you may not use this file except in compliance

 * with the License.  You may obtain a copy of the License at

 *

 *     http://www.apache.org/licenses/LICENSE-2.0

 *

 * Unless required by applicable law or agreed to in writing, software

 * distributed under the License is distributed on an "AS IS" BASIS,

 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

 * See the License for the specific language governing permissions and

 * limitations under the License.

 */

package com.bff.gaia.unified.runners.gaia.translation.wrappers.streaming;



import com.bff.gaia.api.common.state.ListState;

import com.bff.gaia.api.common.state.ListStateDescriptor;

import com.bff.gaia.api.java.functions.KeySelector;

import com.bff.gaia.runtime.state.KeyedStateBackend;

import com.bff.gaia.runtime.state.VoidNamespace;

import com.bff.gaia.runtime.state.VoidNamespaceSerializer;



import java.util.List;

import java.util.stream.Collectors;

import java.util.stream.Stream;

import java.util.stream.StreamSupport;



/**

 * {@link PushedBackElementsHandler} that stores elements in Gaia keyed state, for use when an

 * operation is keyed and pushed-back data needs to stay in the correct partition when they get

 * moved.

 */

class KeyedPushedBackElementsHandler<K, T> implements PushedBackElementsHandler<T> {



  static <K, T> KeyedPushedBackElementsHandler<K, T> create(

      KeySelector<T, K> keySelector,

      KeyedStateBackend<K> backend,

      ListStateDescriptor<T> stateDescriptor) {

    return new KeyedPushedBackElementsHandler<>(keySelector, backend, stateDescriptor);

  }



  private final KeySelector<T, K> keySelector;

  private final KeyedStateBackend<K> backend;

  private final ListStateDescriptor<T> stateDescriptor;



  private KeyedPushedBackElementsHandler(

      KeySelector<T, K> keySelector,

      KeyedStateBackend<K> backend,

      ListStateDescriptor<T> stateDescriptor) {

    this.keySelector = keySelector;

    this.backend = backend;

    this.stateDescriptor = stateDescriptor;

  }



  @Override

  public Stream<T> getElements() {



    return backend

        .getKeys(stateDescriptor.getName(), VoidNamespace.INSTANCE)

        .flatMap(

            key -> {

              try {

                backend.setCurrentKey(key);



                ListState<T> state =

                    backend.getPartitionedState(

                        VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor);



                return StreamSupport.stream(state.get().spliterator(), false);

              } catch (Exception e) {

                throw new RuntimeException("Error reading keyed state.", e);

              }

            });

  }



  @Override

  public void clear() throws Exception {

    // TODO we have to collect all keys because otherwise we get ConcurrentModificationExceptions

    // from gaia. We can change this once it's fixed in Gaia



    List<K> keys =

        backend

            .getKeys(stateDescriptor.getName(), VoidNamespace.INSTANCE)

            .collect(Collectors.toList());



    for (K key : keys) {

      backend.setCurrentKey(key);



      ListState<T> state =

          backend.getPartitionedState(

              VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor);



      state.clear();

    }

  }



  @Override

  public void pushBack(T element) throws Exception {

    ListState<T> state =

        backend.getPartitionedState(

            VoidNamespace.INSTANCE, VoidNamespaceSerializer.INSTANCE, stateDescriptor);



    backend.setCurrentKey(keySelector.getKey(element));

    state.add(element);

  }



  @Override

  public void pushBackAll(Iterable<T> elements) throws Exception {

    for (T e : elements) {

      pushBack(e);

    }

  }

}